import gradio as gr
import numpy as np
import joblib  # 使用 joblib 加载模型
from PIL import Image
import pinecone
from pinecone import Pinecone, ServerlessSpec


# 初始化 Pinecone
api_key = "8426bf47-0774-4cb3-b4b2-7347f267a02e"  # 替换为你的 Pinecone API key
environment = "us-east-1"  # 替换为你的实际环境
pinecone_instance = Pinecone(api_key=api_key)

# 设置 Pinecone 索引
index_name = "mnist-index"
index = pinecone_instance.Index(index_name)

def predict_digit_with_pinecone(image):
    # 检查 image 是否为字典，如果是，则提取图像数据
    if isinstance(image, dict) and "image" in image:
        image = image["image"]
    
    # 确保输入是 NumPy 数组
    image = np.array(image)
    
    # 将输入的手写图像转换为 PIL 格式的灰度图，并调整大小为8x8
    image = Image.fromarray(image).convert('L').resize((8, 8))
    
    # 将图像数据转换为数组，并将其展平为1D数组
    img_array = np.array(image).reshape(1, -1)
    
    # 将像素值缩放到0-1之间
    img_array = (img_array / 255) * 16

    # 准备要查询的向量
    vector = img_array.tolist()[0]
    
    # 使用 Pinecone 查询最近邻（假设已经将训练数据上传到索引）
    query_result = index.query(vector=vector, top_k=1, include_metadata=True)

    # 提取查询结果中的预测标签
    if query_result["matches"]:
        predicted_label = query_result["matches"][0]["metadata"]["label"]
    else:
        predicted_label = "无法识别"

    return int(predicted_label)

# 创建 Gradio 接口
interface = gr.Interface(
    fn=predict_digit_with_pinecone,  # 调用的函数
    inputs=gr.Sketchpad(),  # 手写板输入
    outputs=gr.Label(),  # 输出是预测的标签
    title="手写数字识别",  # Web 应用标题
    description="在手写板上输入一个数字，模型将通过 Pinecone 返回预测结果。"  # 描述信息
)

# 启动 Web 应用
interface.launch()